import pandas as pd
from app_config import get_engine_ts
import numpy as np


def str6(value):
    # 将数字转换为字符串
    str_value = str(value)
    # 如果字符串长度不足6位，则在前面补0
    formatted_value = str_value.zfill(6)
    return formatted_value


def index_weight( from_file, to_file, head_key, three_year_ago, one_year_ago, zero_year_ago):
    engine = get_engine_ts()

    index_930955_index = pd.read_excel(from_file)
    index_930955_index['symbol'] = index_930955_index[head_key].apply(str6)

    query = f""" select symbol, ts_code,industry from stock_basic """
    stock_basic = pd.read_sql(query, get_engine_ts())

    index_930955_index = pd.merge(index_930955_index, stock_basic, on='symbol', how='left')
    # index_930955_index['con_code'] = index_930955_index['ts_code']
    s930955_ts_code = index_930955_index['ts_code']

    sql_dividend = f"""
          SELECT distinct(id) as id ,ts_code, pay_date, cash_div_tax, base_share FROM dividend 
          WHERE div_proc = '实施' AND cash_div_tax > 0 
          AND ts_code IN ({','.join(f"'{code}'" for code in s930955_ts_code)})
          AND pay_date >= '{three_year_ago}' AND pay_date <= '{zero_year_ago}'
          """
    df_dividend_dataset = pd.read_sql_query(sql_dividend, engine)
    df_dividend_dataset['cash_amount'] = df_dividend_dataset['cash_div_tax'] * df_dividend_dataset['base_share']
    df_dividend_tsCode2cashAmount = df_dividend_dataset.groupby('ts_code')['cash_amount'].sum().reset_index()
    df_dividend_tsCode2cashAmount.columns = ['ts_code', 'cash_amount_total']

    query = f"""
          SELECT ts_code,total_mv FROM `daily_basic`
          WHERE ts_code IN ({','.join(f"'{code}'" for code in s930955_ts_code)})
          AND trade_date = '{zero_year_ago}'
          """
    df_total_mv = pd.read_sql_query(query, engine)

    df_temp = pd.merge(index_930955_index, df_dividend_tsCode2cashAmount, on='ts_code', how='left')

    df_dividend_for3year_final = pd.merge(df_temp, df_total_mv, on='ts_code', how='left')

    df_dividend_for3year_final['股息率'] = df_dividend_for3year_final['cash_amount_total'] / 3 / df_dividend_for3year_final['total_mv']

    sql_daily = f"""
           SELECT ts_code, trade_date, pct_chg, amount,close FROM `daily`
           WHERE ts_code IN ({','.join(f"'{code}'" for code in s930955_ts_code)})
           AND trade_date >= '{one_year_ago}'
           AND trade_date <= '{zero_year_ago}'
           """

    df_daily = pd.read_sql_query(sql_daily, engine)

    dict_volatility = {}
    # 遍历计算波动率
    for stock_code, group_df in df_daily.groupby('ts_code'):
        # dict_volatility[stock_code] = calculate_volatility(group_df)
        dict_volatility[stock_code] = group_df.sort_values(by='trade_date')['pct_chg'].std()
    # 将字典转换为 DataFrame
    df_volatility = pd.DataFrame(list(dict_volatility.items()), columns=['ts_code', '波动率'])

    df_dividend_for3year_final_head300 = pd.merge(df_dividend_for3year_final, df_volatility, on='ts_code',
                                                  how='left')
    df_dividend_for3year_final_head300['weight_1'] = df_dividend_for3year_final_head300['股息率'] / df_dividend_for3year_final_head300['波动率']

    df_dividend_for3year_final_head300.to_excel(to_file)


if __name__ == '__main__':
    _from_file = "file/930955cons.xls"
    _to_file = "file/to_930955红利低波100.xlsx"
    _head_key = '成分券代码Constituent Code'
    _three_year_ago = '20220801'
    _one_year_ago = '20240801'
    _zero_year_ago = '20250731'

    index_weight(from_file=_from_file, to_file=_to_file, head_key=_head_key, three_year_ago=_three_year_ago, one_year_ago=_one_year_ago,
                 zero_year_ago=_zero_year_ago)
